#include "vmlinux.h"
#include "cap.bpf.h"
#include <asm/unistd_64.h>

#include <bpf/bpf_helpers.h>
#include <bpf/bpf_core_read.h>

struct s_filter filter = {
	.pid = -1,
	.uid = -1,
};

struct {
	__uint(type, BPF_MAP_TYPE_PERF_EVENT_ARRAY);
	__uint(key_size, sizeof(u32));
	__uint(value_size, sizeof(u32));
}events SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_HASH);
	__type(key, pid_t);
	__type(value, struct s_event);
	__uint(max_entries, 10000);
}rsyscall_enter SEC(".maps");
  
__always_inline static void save_cred(struct task_struct *task, struct s_cap* cap) {
	const struct cred* real_cred = BPF_CORE_READ(task, real_cred);

	cap->securebits = BPF_CORE_READ(real_cred, securebits);

	kernel_cap_t kcap;

	kcap = BPF_CORE_READ(real_cred, cap_inheritable);
	cap->cap_inheritable = ((u64)kcap.cap[1] << 32) + kcap.cap[0];

	kcap = BPF_CORE_READ(real_cred, cap_permitted);
	cap->cap_permitted = ((u64)kcap.cap[1] << 32) + kcap.cap[0];

	kcap = BPF_CORE_READ(real_cred, cap_effective);
	cap->cap_effective = ((u64)kcap.cap[1] << 32) + kcap.cap[0];

	kcap = BPF_CORE_READ(real_cred, cap_bset);
	cap->cap_bset = ((u64)kcap.cap[1] << 32) + kcap.cap[0];

	kcap = BPF_CORE_READ(real_cred, cap_ambient);
	cap->cap_ambient = ((u64)kcap.cap[1] << 32) + kcap.cap[0];
}

SEC("tracepoint/raw_syscalls/sys_enter")
int trace_raw_syscall_sys_enter(struct trace_event_raw_sys_enter *ctx) {
	struct task_struct *task = (struct task_struct *)bpf_get_current_task();
	pid_t pid = BPF_CORE_READ(task, pid);
	pid_t tgid = BPF_CORE_READ(task, tgid);

	uid_t uid = bpf_get_current_uid_gid() >> 32;

	if (filter.pid > 0 && filter.pid != tgid)
		return 0;
	if (filter.uid > 0 && filter.uid != uid)
		return 0;

	pid_t ppid = BPF_CORE_READ(task, real_parent, pid);

	struct s_event event = {
		.nsec =  bpf_ktime_get_ns(),
		.nr = ctx->id,
		.pid = pid,
		.tgid = tgid,
		.ppid = ppid,
		.uid = uid,
	};

	bpf_get_current_comm(event.comm, sizeof(event.comm));
	event.ustack_sz = bpf_get_stack(ctx, event.ustack, sizeof(event.ustack), BPF_F_USER_STACK);

	save_cred(task, &event.cap_before);

	/* 根据系统调用 收集信息 */
	if (ctx->id == __NR_capset) {
		bpf_probe_read_user(&event.data.capset.hdrp, sizeof(struct __user_cap_header_struct), (void*)ctx->args[0]);
		bpf_probe_read_user(&event.data.capset.datap, sizeof(struct __user_cap_data_struct), (void*)ctx->args[1]);
	}
	else if (ctx->id == __NR_execve || ctx->id == __NR_execveat) {
		bpf_probe_read_user(&event.data.execve.filename, sizeof(event.data.execve.filename), (void*)ctx->args[0]);
	}
	else if (ctx->id == __NR_setuid || ctx->id == __NR_setfsuid) {
		event.data.setuid.uid = ctx->args[0];
	}
	else if (ctx->id == __NR_setreuid) {
		event.data.setuid.ruid = ctx->args[0];
		event.data.setuid.euid = ctx->args[1];
	}
	else if (ctx->id == __NR_setresuid) {
		event.data.setuid.ruid = ctx->args[0];
		event.data.setuid.euid = ctx->args[1];
		event.data.setuid.suid = ctx->args[2];
	}
	else if (ctx->id == __NR_setgid || ctx->id == __NR_setfsgid) {
		event.data.setgid.gid = ctx->args[0];
	}
	else if (ctx->id == __NR_setpgid) {
		event.data.setgid.pid = ctx->args[0];
		event.data.setgid.pgid = ctx->args[1];
	}
	else if (ctx->id == __NR_setregid) {
		event.data.setgid.rgid = ctx->args[0];
		event.data.setgid.egid = ctx->args[1];
	}
	else if (ctx->id == __NR_setresgid) {
		event.data.setgid.rgid = ctx->args[0];
		event.data.setgid.egid = ctx->args[1];
		event.data.setgid.sgid = ctx->args[2];
	}
	else if (ctx->id == __NR_clone) {
		event.data.clone.child_pid = ctx->args[0];
	}
	else if (ctx->id == __NR_prctl) {
		event.data.prctl.option = ctx->args[0];
		event.data.prctl.arg2 = ctx->args[1];
		event.data.prctl.arg3 = ctx->args[2];
		event.data.prctl.arg4 = ctx->args[3];
		event.data.prctl.arg5 = ctx->args[4];
	}

	// 第一次跟踪到的输出
	// if (bpf_map_lookup_elem(&rsyscall_enter, &pid) == NULL)
	// 	bpf_perf_event_output(ctx, &events, BPF_F_CURRENT_CPU, &event, sizeof(event));

	bpf_map_update_elem(&rsyscall_enter, &pid, &event, BPF_ANY);

	return 0;
}

SEC("tracepoint/raw_syscalls/sys_exit")
int trace_raw_syscall_sys_exit(struct trace_event_raw_sys_exit *ctx) {
	struct task_struct *task = (struct task_struct *)bpf_get_current_task();
	pid_t pid = BPF_CORE_READ(task, pid);
	uid_t uid = bpf_get_current_uid_gid() >> 32;

	if (filter.pid > 0 && filter.pid != pid)
		return 0;
	if (filter.uid > 0 && filter.uid != uid)
		return 0;

	struct s_event* event = bpf_map_lookup_elem(&rsyscall_enter, &pid);
	if (event == NULL)
		return 0;

	pid_t ppid = BPF_CORE_READ(task, real_parent, pid);
	pid_t tgid = BPF_CORE_READ(task, tgid);
	event->ppid = ppid;
	event->tgid = tgid;

	struct s_cap new_cap;
	save_cred(task, &new_cap);

	if (new_cap.cap_inheritable != event->cap_before.cap_inheritable
		|| new_cap.cap_permitted != event->cap_before.cap_permitted
		|| new_cap.cap_effective != event->cap_before.cap_effective
		|| new_cap.cap_bset != event->cap_before.cap_bset
		|| new_cap.cap_ambient != event->cap_before.cap_ambient
		|| new_cap.securebits != event->cap_before.securebits)
	{
		event->cap_after = new_cap;
		event->nsec = bpf_ktime_get_ns();
		bpf_map_update_elem(&rsyscall_enter, &pid, event, BPF_EXIST);

		event->retvel = ctx->ret;
		bpf_perf_event_output(ctx, &events, BPF_F_CURRENT_CPU, event, sizeof(*event));
	}

	return 0;
}

char _license[] SEC("license") = "GPL";

__always_inline static int report_maps_event(void* ctx, int nr) {
	pid_t pid = bpf_get_current_pid_tgid() & __UINT32_MAX__;
	pid_t tgid = bpf_get_current_pid_tgid() >> 32;
	uid_t uid = bpf_get_current_uid_gid() & __UINT32_MAX__;

	struct task_struct *task = (struct task_struct *)bpf_get_current_task();
	pid_t ppid = BPF_CORE_READ(task, real_parent, pid);

	if (filter.pid > 0 && filter.pid != tgid)
		return 0;
	if (filter.uid > 0 && filter.uid != uid)
		return 0;

	struct s_maps_event event = {
		.pid = pid,
		.tgid = tgid,
		.ppid = ppid,
		.nr = nr,
	};

	bpf_perf_event_output(ctx, &events, BPF_F_CURRENT_CPU, &event, sizeof(event));

	return 0;
}

SEC("kprobe/do_exit")
int kprobe_do_exit(struct trace_kprobe *ctx) {
	pid_t pid = bpf_get_current_pid_tgid() & __UINT32_MAX__;
	bpf_map_delete_elem(&rsyscall_enter, &pid);

	return report_maps_event(ctx, -1);
}